import os
import time
import base64
import shutil
import random
import waitress
import zipstream
import filesystem as fs
from   werkzeug.security       import generate_password_hash, check_password_hash
from   werkzeug.datastructures import Headers
from   urllib.parse import unquote
from   functools    import wraps
from   typing       import Generator
from   flask        import Flask, render_template, request, session, Response, stream_with_context

USERS = {
    'Administrator': generate_password_hash('admin123')
}

def calculate_zip_size(folder_path:str, is_zip64:bool=True ) -> int:
    if  is_zip64 and calculate_zip_size(folder_path,  False) < 2*1024*1024*1024:
        is_zip64 = False
    total_size = 0
    file_count = 0
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                file_size     = os.path.getsize(file_path)
                relative_path = os.path.relpath(file_path, folder_path)
                filename_len  = len(relative_path.encode('utf-8'))
                total_size += filename_len * 2 + file_size + 120 + (12 if is_zip64 else 0)
                file_count += 1
            except (OSError, IOError):return -1
    return total_size if file_count == 0 else total_size + 22 + (52 if is_zip64 else 0)

def get_modification_time(path) -> str:
    try:
        if not os.path.exists(path):
            return format_time(time.time())
        return format_time(os.path.getmtime(path))
    except (OSError, PermissionError):
        return format_time(time.time())
    
def format_time(timestamp:float, mode:str='%a, %d %b %Y %H:%M:%S GMT') -> str:
    return time.strftime(mode, time.localtime(timestamp))

def stream_file(file_path: str, chunk_size: int = 5*1024*1024) -> Generator[bytes, None, None]:
    with open(file_path, 'rb') as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:break
            yield  chunk

def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        # 如果已经登录，直接放行
        if session.get('logged_in', ''):
            return f(*args, **kwargs)
            
        # 如果没有用户配置，允许匿名访问
        if not USERS:
            session['logged_in'] = True
            session['username'] = 'anonymous'
            return f(*args, **kwargs)
            
        # 检查基本认证
        auth = request.authorization
        if auth and auth.username in USERS and check_password_hash(USERS[auth.username], auth.password):
            # 认证成功，设置session
            session['logged_in'] = True
            session['username'] = auth.username
            return f(*args, **kwargs)
            
        # 需要认证
        return Response(
            'Authentication required', 401,
            {'WWW-Authenticate': 'Basic realm="WebDAV Login"'}
        )
    return decorated

def create_random_key(length:int=32) -> str:
    return '^'.join(random.choice('0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') for i in range(length))

app = Flask(__name__)
app.secret_key = create_random_key(2048)

@app.route('/', methods=['LOGOUT'])
@app.route('/<path:path>', methods=['LOGOUT'])
def logout(path='/'):
    session.pop('logged_in', None)
    session.pop('username', None)
    return 'Logged out', 200

@app.route('/', methods=['PROPFIND'])
@app.route('/<path:path>', methods=['PROPFIND'])
@requires_auth
def list_directory(path='/'):
    def generate(href, name, is_dir, size, modified):
        request_data  = '<D:response>\n<D:href>{}</D:href>\n'.format(href)
        request_data += '<D:propstat>\n<D:prop>\n'
        request_data += '<D:displayname>{}</D:displayname>\n'.format(name)
        request_data += '<D:getcontentlength>{}</D:getcontentlength>\n'.format(size) if not is_dir else ''
        request_data += '<D:getlastmodified>{}</D:getlastmodified>\n'.format(modified)
        request_data += '<D:resourcetype>{}</D:resourcetype>'.format('<D:collection/>' if is_dir else '')
        request_data += '\n</D:prop>\n<D:status>HTTP/1.1 200 OK</D:status>\n</D:propstat>\n</D:response>'
        return request_data
    def get_info(path:str) -> tuple[str, str, bool, int, str]: # returns (href, name, is_dir, size, modified)
        if  not  path:
            return '/', '', True, 0, time.strftime('%a, %d %b %Y %H:%M:%S GMT', format_time(time.time()))
        real_path = fs.real_path(path)
        is_dir = path.endswith('/')
        try   :size = os.path.getsize(real_path) if not is_dir else 0
        except:size = 0
        return path, path.split('/')[-1] if not is_dir else path.split('/')[-2], is_dir, size, get_modification_time(real_path) 
    path = unquote(path if path.endswith('/') else path + '/')
    path = path if path.startswith('/') else '/' + path
    if  not os.path.isdir(fs.real_path(path)) and not os.path.isfile(fs.real_path(path[:-1])) and path != '/':
        return 'Not Found', 404
    if  os.path.isfile(fs.real_path(path[:-1])):
        request_data = '<?xml version="1.0" encoding="utf-8" ?>\n<D:multistatus xmlns:D="DAV:">\n' + generate(*get_info(path[:-1]))
        return request_data + '</D:multistatus>', 207
    List = fs.list_directory(path)[0]
    request_data = '<?xml version="1.0" encoding="utf-8" ?>\n<D:multistatus xmlns:D="DAV:">\n' + generate(*get_info(path))
    for file in List:
        request_data += generate(*get_info(path + file))
    return request_data + '</D:multistatus>', 207

@app.route('/', methods=['PROPPATCH'])
@app.route('/<path:path>', methods=['PROPPATCH'])
@requires_auth
def proppatch_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path or not os.path.isfile(real_path):
        return 'Not Found', 404
    request_data  = '<?xml version="1.0" encoding="utf-8" ?>\n'
    request_data += '<D:multistatus xmlns:D="DAV:">\n'
    request_data += '<D:response>\n<D:href>{}</D:href>\n'.format(path)
    request_data += '<D:propstat>\n<D:prop>\n</D:prop>\n'
    request_data += '<D:status>HTTP/1.1 403 Forbidden</D:status>\n'
    request_data += '</D:propstat>\n</D:response>\n</D:multistatus>'
    return request_data, 207

@app.route('/', methods=['OPTIONS'])
def options():
    headers = {
        'DAV': '1, 2',
        'Allow': 'OPTIONS, GET, HEAD, POST, PUT, DELETE, PROPFIND, PROPPATCH, MKCOL, COPY, MOVE, LOCK, UNLOCK',
        'MS-Author-Via': 'DAV'
    }
    return Response(status=200, headers=headers)

@app.route('/', methods=['HEAD'])
@app.route('/<path:path>', methods=['HEAD'])
@requires_auth
def head_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return Response(status=405)
    if not os.path.exists(real_path):
        return Response(status=404)
    return Response(status=200)

@app.route('/static/<path:path>', methods=['GET'])
def get_static(path):
    return app.send_static_file(path)

@app.route('/', methods=['GET'])  #  GET Index Page
@app.route('/<path:path>', methods=['GET'])
def get_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    Authorization = request.headers.get('Authorization', '')
    if 'Basic' in Authorization and not session.get('logged_in', ''):
        username, password = base64.b64decode(Authorization.split(' ')[1]).decode().split(':')
        if  username in USERS and check_password_hash(USERS[username], password):
            session['logged_in'] = True
            session['username'] = username
        else:return 'Invalid username or password', 401
    if os.path.isfile(real_path):
        if  not session.get('logged_in', ''):
            return Response(
                'Authentication required', 401,
                {'WWW-Authenticate': 'Basic realm="WebDAV Login"'}
            )
        try:
            file_size = os.path.getsize(real_path)
            modified_time = get_modification_time(real_path)
            
            # 检查Range请求头
            range_header = request.headers.get('Range', None)
            if range_header:
                # 解析Range头
                byte1, byte2 = 0, None
                range_units = range_header.replace('bytes=', '').split('-')
                byte1 = int(range_units[0])
                if range_units[1]:
                    byte2 = int(range_units[1])
                
                # 计算实际要发送的字节范围
                if byte2 is None:
                    byte2 = file_size - 1
                else:
                    byte2 = min(byte2, file_size - 1)
                
                length = byte2 - byte1 + 1
                
                # 设置部分内容响应头
                headers = Headers()
                headers.add('Content-Type', 'application/octet-stream')
                headers.add('Content-Length', str(length))
                headers.add('Last-Modified', modified_time)
                headers.add('Accept-Ranges', 'bytes')
                headers.add('Content-Range', f'bytes {byte1}-{byte2}/{file_size}')
                
                # 生成文件片段流
                def generate_file_chunk():
                    with open(real_path, 'rb') as f:
                        f.seek(byte1)
                        remaining = length
                        while remaining > 0:
                            chunk_size = min(8192, remaining)
                            chunk = f.read(chunk_size)
                            if not chunk:
                                break
                            remaining -= len(chunk)
                            yield chunk
                
                return Response(
                    generate_file_chunk(),
                    status=206,  # Partial Content
                    headers=headers
                )

            # 设置响应头
            headers = Headers()
            headers.add('Content-Type', 'application/octet-stream')
            headers.add('Content-Length', str(file_size))
            headers.add('Last-Modified', modified_time)
            headers.add('Accept-Ranges', 'bytes')

            return Response(
                stream_with_context(stream_file(real_path)),
                headers=headers,
                status=200
            ) if file_size > 20*1024*1024 else Response(
                open(real_path, 'rb').read(),
                headers=headers,
                status=200
            )
        except (OSError, PermissionError):
            return '', 405
    if  session.get('logged_in', None):
        request_args = request.args.to_dict()
        if 'simple' in request_args:
            List = fs.list_directory(path)[0]
            return '\n'.join(List), 200
        if 'zip' in request_args:
            def generate_zip_file() -> Generator[bytes, None, None]:
                z = zipstream.ZipStream()

                def process_directory(dir_path, base_arcname=''):
                    try:
                        items = os.listdir(dir_path)
                        for item in items:
                            item_path = os.path.join(dir_path, item)
                            arcname = os.path.join(base_arcname, item).replace('\\', '/') if base_arcname else item
                            if os.path.isfile(item_path):
                                z.add(open(item_path, 'rb'), arcname=arcname, compress_type=zipstream.ZIP_STORED)
                            if os.path.isdir(item_path):
                                process_directory(item_path, arcname)
                    except (OSError, PermissionError):pass

                process_directory(real_path)

                for chunk in z:
                    yield chunk

            headers = Headers()
            headers.add('Content-Type', 'application/zip')
            headers.add('Content-Disposition', f'attachment; filename="{path.split("/")[-2].encode('utf-8').decode('latin-1')}.zip"')
            headers.add('Last-Modified', get_modification_time(real_path))
            headers.add('Content-Length', str(calculate_zip_size(real_path)))
            return Response(
                generate_zip_file(),
                headers=headers,
                status=200
            )
        
        if 'size' in request_args:
            return str(calculate_zip_size(real_path)), 200
        return render_template('index.html', username=session.get('username', ''))
    return render_template('login.html')
    
@app.route('/', methods=['PUT'])
@app.route('/<path:path>', methods=['PUT'])
@requires_auth
def put_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return '', 405
    if os.path.isdir(real_path):
        return '', 405
    if not os.path.isdir(os.path.dirname(real_path)):
        os.makedirs(os.path.dirname(real_path), exist_ok=True)
    try:
        with open(real_path, 'wb') as f:
            while True:
                chunk = request.stream.read(10 * 1024 * 1024)
                if not  chunk:break
                f.write(chunk)
            f.flush()
    except (OSError, PermissionError) as Error:
        return 'Error: ' + Error, 500
    return '', 201

@app.route('/', methods=['MOVE'])
@app.route('/<path:path>', methods=['MOVE'])
@requires_auth
def move_directory(path='/'):
    def get_dst_path() -> str:
        dst_path = unquote(request.headers.get('Destination'))
        if not dst_path:return ''
        if 'http://'  in dst_path:dst_path = dst_path.replace('http://' , '')
        if 'https://' in dst_path:dst_path = dst_path.replace('https://', '')
        return fs.real_path('/' + '/'.join(dst_path.split('/')[1:]))
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return '', 405
    if not os.path.exists(real_path):
        return '', 404
    dst_path = request.headers.get('Destination')
    if not dst_path:
        return '', 400
    dst_real_path = get_dst_path()
    if not dst_real_path:
        return '', 405
    if not os.path.exists(os.path.dirname(dst_real_path)):
        if os.path.isfile(real_path):
            os.makedirs(os.path.dirname(dst_real_path), exist_ok=True)
    try:
        shutil.move(real_path, dst_real_path)
        return '', 201
    except (shutil.Error, OSError, PermissionError):
        return '', 500

@app.route('/', methods=['COPY'])
@app.route('/<path:path>', methods=['COPY'])
@requires_auth
def copy_directory(path='/'):
    def get_dst_path() -> str:
        dst_path = unquote(request.headers.get('Destination'))
        if not dst_path:return ''
        if 'http://'  in dst_path:dst_path = dst_path.replace('http://' , '')
        if 'https://' in dst_path:dst_path = dst_path.replace('https://', '')
        return fs.real_path('/' + '/'.join(dst_path.split('/')[1:]))
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return '', 405
    if not os.path.exists(real_path):
        return '', 404
    dst_real_path = get_dst_path()
    if not dst_real_path:
        return '', 405
    if not os.path.exists(os.path.dirname(dst_real_path)):
        if os.path.isfile(real_path):
            os.makedirs(os.path.dirname(dst_real_path), exist_ok=True)
    try:
        if os.path.isfile(real_path):
            shutil.copy(real_path, dst_real_path)
        else:
            shutil.copytree(real_path, dst_real_path, dirs_exist_ok=True)
        return '', 201
    except (shutil.Error, OSError, PermissionError):
        return '', 500

@app.route('/', methods=['MKCOL'])
@app.route('/<path:path>', methods=['MKCOL'])
@requires_auth
def mkcol_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return '', 405
    if os.path.exists(real_path):
        return '', 405
    try:
        os.makedirs(real_path)
        return '', 201
    except (OSError, PermissionError):
        return '', 500

@app.route('/', methods=['DELETE'])
@app.route('/<path:path>', methods=['DELETE'])
@requires_auth
def delete_directory(path='/'):
    path = unquote(path)
    path = path if path.startswith('/') else '/' + path
    real_path = fs.real_path(path)
    if not real_path:
        return '', 405
    if not os.path.exists(real_path):
        return '', 404
    if os.path.isfile(real_path):
        try:
            os.remove(real_path)
            return '', 204
        except (OSError, PermissionError):
            return '', 500
    if os.path.isdir(real_path):
        try:
            shutil.rmtree(real_path)
            return '', 204
        except (shutil.Error, OSError, PermissionError):
            return '', 500

if __name__ == '__main__':
    waitress.serve(app, host='0.0.0.0', port=5002, threads=64, max_request_body_size=1e10, connection_limit=1e10, channel_timeout=1e10)
    # app.run(host='0.0.0.0', port=5002, debug=True)